# tools/mcp_import_manager.py
# MCP服务器导入管理器，支持从外部JSON配置导入MCP工具

import json
import shutil
from pathlib import Path
from typing import Dict, List, Any, Optional
from datetime import datetime
import asyncio
import logging
import traceback

from .mcp_client import MCPClientManager


class MCPServerImportManager:
    """MCP服务器导入管理器，负责从外部配置导入MCP工具"""
    
    def __init__(self, config_path: str = "configs/mcp_server.json"):
        self.config_path = Path(config_path)
        self.backup_dir = Path("configs/backups")
        self.logger = self._setup_logger()
        self.client_managers: Dict[str, MCPClientManager] = {}
        
    def _setup_logger(self) -> logging.Logger:
        """设置日志记录器"""
        logger = logging.getLogger("mcp_import_manager")
        logger.setLevel(logging.DEBUG)
        
        if not logger.handlers:
            # 控制台处理器
            console_handler = logging.StreamHandler()
            console_formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            console_handler.setFormatter(console_formatter)
            console_handler.setLevel(logging.INFO)
            logger.addHandler(console_handler)
            
            # 文件处理器 - 详细日志
            log_file = Path("mcp_debug.log")
            file_handler = logging.FileHandler(log_file, encoding='utf-8')
            file_formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(funcName)s:%(lineno)d - %(message)s'
            )
            file_handler.setFormatter(file_formatter)
            file_handler.setLevel(logging.DEBUG)
            logger.addHandler(file_handler)
            
        return logger
    
    def load_config(self) -> Dict[str, Any]:
        """加载MCP服务器配置"""
        try:
            if not self.config_path.exists():
                self.logger.warning(f"配置文件不存在: {self.config_path}")
                return {}
                
            with open(self.config_path, 'r', encoding='utf-8') as f:
                config = json.load(f)
                self.logger.info(f"成功加载配置文件: {self.config_path}")
                return config
        except Exception as e:
            self.logger.error(f"加载配置文件失败: {e}")
            return {}
    
    def save_config(self, config: Dict[str, Any]) -> bool:
        """保存MCP服务器配置"""
        try:
            # 创建备份
            if config.get("importSettings", {}).get("backupBeforeImport", True):
                self._create_backup()
            
            # 更新最后修改时间
            config["lastUpdated"] = datetime.now().isoformat() + "Z"
            
            with open(self.config_path, 'w', encoding='utf-8') as f:
                json.dump(config, f, indent=2, ensure_ascii=False)
                
            self.logger.info(f"成功保存配置文件: {self.config_path}")
            return True
        except Exception as e:
            self.logger.error(f"保存配置文件失败: {e}")
            return False
    
    def _create_backup(self) -> None:
        """创建配置文件备份"""
        try:
            self.backup_dir.mkdir(exist_ok=True)
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            backup_path = self.backup_dir / f"mcp_server_{timestamp}.json"
            shutil.copy2(self.config_path, backup_path)
            self.logger.info(f"创建备份文件: {backup_path}")
        except Exception as e:
            self.logger.warning(f"创建备份失败: {e}")
    
    def validate_config(self, config: Dict[str, Any]) -> List[str]:
        """验证配置文件 - 简化版"""
        errors = []
        
        if "mcpServers" not in config or not isinstance(config["mcpServers"], dict):
            errors.append("缺少有效的 'mcpServers' 配置")
            return errors
        
        for server_id, server_config in config["mcpServers"].items():
            if not isinstance(server_config, dict):
                errors.append(f"服务器 '{server_id}' 配置无效")
                continue
            
            server_type = self._detect_server_type(server_config)
            # 检查必需字段
            if server_type == "stdio" and "command" not in server_config:
                errors.append(f"stdio服务器 '{server_id}' 缺少 'command'")
            elif server_type in ["sse", "http", "streamableHttp"] and "url" not in server_config:
                errors.append(f"{server_type}服务器 '{server_id}' 缺少 'url'")
        
        return errors
    
    def _detect_server_type(self, server_config: Dict[str, Any]) -> str:
        """检测服务器类型"""
        # 如果明确指定了type字段
        if "type" in server_config:
            return server_config["type"]
        
        # 如果有command字段，默认为stdio
        if "command" in server_config:
            return "stdio"
        
        # 如果有url字段，默认为sse
        if "url" in server_config:
            return "sse"
        
        return "unknown"
    
    def import_from_json(self, json_data: str) -> Dict[str, Any]:
        """从JSON导入配置 - 简化版"""
        try:
            imported_config = json.loads(json_data)
            
            # 验证配置
            errors = self.validate_config(imported_config)
            if errors:
                return {"success": False, "errors": errors, "message": "配置验证失败"}
            
            # 合并配置
            existing_config = self.load_config()
            if "mcpServers" not in existing_config:
                existing_config["mcpServers"] = {}
            
            for server_id, server_config in imported_config.get("mcpServers", {}).items():
                existing_config["mcpServers"][server_id] = server_config
                self.logger.info(f"导入服务器: {server_id}")
            
            # 保存配置
            if self.save_config(existing_config):
                return {
                    "success": True,
                    "message": "配置导入成功",
                    "imported_servers": list(imported_config.get("mcpServers", {}).keys())
                }
            return {"success": False, "message": "保存失败"}
                
        except json.JSONDecodeError as e:
            return {"success": False, "message": f"JSON错误: {e}"}
        except Exception as e:
            return {"success": False, "message": f"导入失败: {e}"}
    
    def get_enabled_servers(self) -> Dict[str, Any]:
        """获取启用的MCP服务器列表"""
        config = self.load_config()
        enabled_servers = {}
        
        for server_id, server_config in config.get("mcpServers", {}).items():
            if server_config.get("enabled", False):
                enabled_servers[server_id] = server_config
        
        return enabled_servers
    
    def enable_server(self, server_id: str) -> bool:
        """启用指定的MCP服务器"""
        config = self.load_config()
        
        if server_id in config.get("mcpServers", {}):
            config["mcpServers"][server_id]["enabled"] = True
            return self.save_config(config)
        return False
    
    def disable_server(self, server_id: str) -> bool:
        """禁用指定的MCP服务器"""
        config = self.load_config()
        
        if server_id in config.get("mcpServers", {}):
            config["mcpServers"][server_id]["enabled"] = False
            return self.save_config(config)
        return False
    
    def remove_server(self, server_id: str) -> bool:
        """移除指定的MCP服务器"""
        config = self.load_config()
        
        if server_id in config.get("mcpServers", {}):
            del config["mcpServers"][server_id]
            return self.save_config(config)
        return False
    
    async def initialize_enabled_servers(self) -> Dict[str, Any]:
        """初始化启用的MCP服务器 - 改进错误处理"""
        enabled_servers = self.get_enabled_servers()
        results = {}
        
        self.logger.info(f"开始初始化 {len(enabled_servers)} 个服务器")
        print(f"\n{'='*60}")
        print(f"📡 开始初始化 {len(enabled_servers)} 个MCP服务器")
        print(f"{'='*60}")
        
        for server_id, server_config in enabled_servers.items():
            try:
                server_type = self._detect_server_type(server_config)
                self.logger.info(f"初始化 {server_id} ({server_type})")
                print(f"\n🔧 初始化 {server_id} ({server_type})...")
                
                client_manager = MCPClientManager()
                success, tools = False, []
                
                if server_type == "stdio":
                    success = await self._initialize_stdio_server(client_manager, server_config)
                    if success:
                        try:
                            tools = await client_manager.list_tools()
                            print(f"✅ {server_id} 初始化成功，获取到 {len(tools)} 个工具")
                        except Exception as e:
                            self.logger.warning(f"获取工具列表失败: {e}")
                            print(f"⚠️  警告: {server_id} 初始化成功，但获取工具列表失败: {e}")
                            tools = []
                elif server_type == "streamableHttp":
                    success = await self._initialize_streamable_http_server(client_manager, server_config)
                    if success:
                        tools = await self._get_streamable_http_tools(server_config)
                        print(f"✅ {server_id} 初始化成功，获取到 {len(tools)} 个工具")
                else:
                    error_msg = f"不支持的类型: {server_type}"
                    self.logger.error(f"× {server_id}: {error_msg}")
                    print(f"❌ {server_id}: {error_msg}")
                    results[server_id] = {"success": False, "error": error_msg}
                    continue
                
                if success:
                    self.client_managers[server_id] = client_manager
                    results[server_id] = {"success": True, "tools": tools}
                    self.logger.info(f"✓ {server_id}: {len(tools)} 个工具")
                else:
                    error_msg = "初始化失败（详见上方错误信息）"
                    results[server_id] = {"success": False, "error": error_msg}
                    print(f"❌ {server_id}: {error_msg}")
                    
            except Exception as e:
                error_msg = str(e)
                self.logger.error(f"× {server_id}: {error_msg}\n{traceback.format_exc()}")
                print(f"❌ {server_id}: 发生异常 - {error_msg}")
                results[server_id] = {"success": False, "error": error_msg}
        
        # 输出总结
        successful = sum(1 for r in results.values() if r.get("success"))
        failed = len(enabled_servers) - successful
        self.logger.info(f"初始化完成: {successful}/{len(enabled_servers)} 成功")
        print(f"\n{'='*60}")
        print(f"📊 初始化完成: {successful}/{len(enabled_servers)} 成功")
        if failed > 0:
            print(f"⚠️  有 {failed} 个服务器初始化失败，但不影响其他功能")
        print(f"{'='*60}\n")
        
        return results
    
    async def _initialize_stdio_server(self, client_manager: MCPClientManager, server_config: Dict[str, Any]) -> bool:
        """初始化stdio服务器 - 改进错误处理"""
        try:
            command = server_config.get("command")
            args = server_config.get("args", [])
            env = server_config.get("env", {})
            
            if not command:
                self.logger.error("缺少command")
                print(f"❌ 错误: 服务器配置缺少 'command' 字段")
                return False
            
            # 检查命令是否可用（仅检查常见命令）
            if command in ["npx", "node", "python"]:
                import shutil
                command_path = shutil.which(command)
                if not command_path:
                    error_msg = f"命令 '{command}' 未找到，请确保已安装并添加到PATH"
                    self.logger.error(error_msg)
                    print(f"❌ 错误: {error_msg}")
                    return False
            
            # 检查文件路径（仅对Python脚本）
            if args and args[0].endswith('.py'):
                file_path = Path(args[0])
                if not file_path.exists():
                    relative_path = Path(".") / args[0]
                    if relative_path.exists():
                        args[0] = str(relative_path)
                    else:
                        error_msg = f"文件不存在: {args[0]}"
                        self.logger.error(error_msg)
                        print(f"❌ 错误: {error_msg}")
                        return False
            
            # 尝试初始化，增加超时保护
            import asyncio
            try:
                result = await asyncio.wait_for(
                    client_manager.initialize_with_command(command, args, env),
                    timeout=10.0
                )
                if not result:
                    self.logger.warning(f"初始化返回False，可能连接失败")
                    print(f"⚠️  警告: 服务器初始化返回失败，可能是连接问题")
                return result
            except asyncio.TimeoutError:
                error_msg = f"初始化超时（>10秒），可能是服务器启动失败或命令错误"
                self.logger.error(error_msg)
                print(f"❌ 错误: {error_msg}")
                return False
            
        except Exception as e:
            error_msg = f"stdio初始化失败: {e}"
            self.logger.error(f"{error_msg}\n{traceback.format_exc()}")
            print(f"❌ 错误: {error_msg}")
            return False
    
    async def call_tool(self, server_id: str, tool_name: str, **kwargs) -> str:
        """调用工具 - 简化版"""
        try:
            config = self.load_config()
            server_config = config.get("mcpServers", {}).get(server_id, {})
            server_type = self._detect_server_type(server_config)
            
            if server_type == "streamableHttp":
                return await self._call_streamable_http_tool(server_config, tool_name, kwargs)
            elif server_id in self.client_managers:
                return await self.client_managers[server_id].call_tool(tool_name, kwargs)
            else:
                return f"服务器 {server_id} 未初始化"
                
        except Exception as e:
            error = f"工具调用失败: {str(e)}"
            self.logger.error(f"{error}\n{traceback.format_exc()}")
            return error
    
    async def close_all_connections(self):
        """关闭所有MCP连接 - 修复异步错误"""
        if not self.client_managers:
            return
        
        connections_to_close = list(self.client_managers.items())
        self.client_managers.clear()
        
        # 顺序关闭连接，避免异步取消作用域错误
        for server_id, client_manager in connections_to_close:
            try:
                await asyncio.wait_for(
                    self._close_single_connection(server_id, client_manager),
                    timeout=3.0
                )
            except asyncio.TimeoutError:
                self.logger.warning(f"关闭 {server_id} 超时")
            except Exception as e:
                self.logger.warning(f"关闭 {server_id} 出错: {e}")
        
        self.logger.info("所有MCP连接已关闭")
    
    async def _close_single_connection(self, server_id: str, client_manager):
        """关闭单个连接 - 添加异常保护"""
        try:
            if hasattr(client_manager, 'close'):
                await client_manager.close()
                self.logger.info(f"✓ 关闭连接: {server_id}")
        except asyncio.CancelledError:
            # 忽略取消错误
            pass
        except Exception as e:
            self.logger.warning(f"关闭 {server_id} 失败: {e}")
    
    async def _initialize_streamable_http_server(self, client_manager: MCPClientManager, server_config: Dict[str, Any]) -> bool:
        """初始化streamableHttp类型的MCP服务器"""
        url = server_config.get("url")
        if not url:
            self.logger.error("streamableHttp服务器缺少url字段")
            return False
        
        self.logger.info(f"初始化streamableHttp服务器: {url}")
        
        try:
            from mcp.client.streamable_http import streamablehttp_client
            from mcp.client.session import ClientSession
            
            client = streamablehttp_client(url)
            read, write, get_session_id = await client.__aenter__()
            session = ClientSession(read, write)
            await session.__aenter__()
            await session.initialize()
            
            # 保存到client_manager
            client_manager.session = session
            client_manager.client = client
            client_manager.read = read
            client_manager.write = write
            
            self.logger.info("streamableHttp连接成功")
            return True
            
        except (ImportError, Exception) as e:
            self.logger.warning(f"MCP客户端初始化失败，回退到HTTP: {e}")
            return True  # 回退到HTTP请求方式
    
    async def _get_streamable_http_tools(self, server_config: Dict[str, Any]) -> List[Dict[str, Any]]:
        """获取streamableHttp服务器的工具列表"""
        try:
            url = server_config.get("url")
            if not url:
                return []
            
            # 尝试使用MCP客户端
            try:
                from mcp.client.streamable_http import streamablehttp_client
                from mcp.client.session import ClientSession
                
                async with streamablehttp_client(url) as (read, write, get_session_id):
                    session = ClientSession(read, write)
                    await session.__aenter__()
                    await session.initialize()
                    
                    tools_result = await session.list_tools()
                    await session.__aexit__(None, None, None)
                    
                    self.logger.info(f"通过MCP客户端获取 {len(tools_result.tools)} 个工具")
                    return tools_result.tools
                
            except (ImportError, Exception) as e:
                self.logger.warning(f"MCP客户端获取工具失败，回退到HTTP: {e}")
            
            # 回退到HTTP请求
            import httpx
            
            mcp_request = {
                "jsonrpc": "2.0",
                "id": 1,
                "method": "tools/list",
                "params": {}
            }
            
            headers = {
                "Content-Type": "application/json",
                "Accept": "application/json, text/event-stream"
            }
            
            async with httpx.AsyncClient(timeout=10.0) as client:
                response = await client.post(url, json=mcp_request, headers=headers)
                response.raise_for_status()
                result = response.json()
                
                if "result" in result and "tools" in result["result"]:
                    tools = result["result"]["tools"]
                    self.logger.info(f"通过HTTP获取 {len(tools)} 个工具")
                    return tools
                
                return []
                    
        except Exception as e:
            self.logger.error(f"获取工具列表失败: {e}")
            return []
    
    async def _call_streamable_http_tool(self, server_config: Dict[str, Any], tool_name: str, kwargs: Dict[str, Any]) -> str:
        """调用streamableHttp服务器的工具 - 优化版"""
        self.logger.debug(f"调用工具: {tool_name}, 参数: {kwargs}")
        
        try:
            url = server_config.get("url")
            if not url:
                return "错误：服务器URL未配置"
            
            # 确保kwargs是字典类型
            if isinstance(kwargs, str):
                import json
                try:
                    kwargs = json.loads(kwargs)
                except json.JSONDecodeError as e:
                    return f"参数格式错误: {e}"
            
            if not isinstance(kwargs, dict):
                kwargs = {}
            
            # 优先使用MCP官方客户端
            try:
                from mcp.client.streamable_http import streamablehttp_client
                from mcp.client.session import ClientSession
                
                async with streamablehttp_client(url) as (read, write, get_session_id):
                    session = ClientSession(read, write)
                    await session.__aenter__()
                    await session.initialize()
                    
                    result = await session.call_tool(tool_name, arguments=kwargs)
                    await session.__aexit__(None, None, None)
                    
                    if result.content and len(result.content) > 0:
                        return result.content[0].text
                    return "工具执行成功，但无返回内容"
                    
            except (ImportError, Exception) as e:
                self.logger.warning(f"MCP客户端调用失败，回退到HTTP: {e}")
            
            # 回退到HTTP请求
            import httpx
            import json
            
            mcp_request = {
                "jsonrpc": "2.0",
                "id": 1,
                "method": "tools/call",
                "params": {
                    "name": tool_name,
                    "arguments": kwargs
                }
            }
            
            headers = {
                "Content-Type": "application/json",
                "Accept": "application/json, text/event-stream"
            }
            
            async with httpx.AsyncClient(timeout=30.0) as client:
                response = await client.post(url, json=mcp_request, headers=headers)
                response.raise_for_status()
                result = response.json()
                
                # 解析响应
                if "result" in result and "content" in result["result"]:
                    content = result["result"]["content"]
                    if content:
                        first_content = content[0]
                        return first_content.get("text", str(first_content))
                    return "工具执行成功，但无返回内容"
                elif "error" in result:
                    return f"工具调用失败: {result['error'].get('message', '未知错误')}"
                
                return "工具调用响应格式异常"
                    
        except Exception as e:
            error_msg = f"调用工具失败: {e}"
            self.logger.error(error_msg)
            return error_msg


# 全局导入管理器实例
mcp_import_manager = MCPServerImportManager()


def import_mcp_servers_from_json(json_data: str) -> Dict[str, Any]:
    """从JSON导入MCP服务器的便捷函数"""
    return mcp_import_manager.import_from_json(json_data)


async def initialize_imported_servers() -> Dict[str, Any]:
    """初始化导入的MCP服务器的便捷函数"""
    return await mcp_import_manager.initialize_enabled_servers()
